import re
import os
import pandas as pd
import numpy as np
import plotly.graph_objs as go
from plotly.subplots import make_subplots
from datetime import timedelta, datetime
from loguru import logger as log
import program.utils.tools as tools
import program.utils.reader as reader
from program.function import filter_generate, filter_before
from program.environ import tag


pd.set_option('expand_frame_repr', False)
pd.set_option('display.max_rows', 1000)  # 最多显示数据的行数

root_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
# 'data/pickle_data/swap'
pickle_path = os.path.join(root_path, 'data', 'pickle_data', 'swap')
# 'data/factors/swap'
factor_pickle = os.path.join(root_path, 'data', 'factors', 'swap')

# 基础配置
hold_hour_num = 1  # hold_hour (get_coins_with_filter_before函数必需)
long_white_list = []  # 固定白名单(如'BTCUSDT')
short_white_list = []  # 固定白名单
long_black_list = []  # 固定黑名单
short_black_list = []  # 固定黑名单


def get_symbols_with_filter_before(_start_date, _end_date, _filter_before_params):
    """
    通过过滤条件，获取某个时间段内过滤后的每个时刻的币种，返回一个dataframe，以时间为索引，包含对应时间的过滤后的币种
    :param _start_date: 分析开始时间
    :param _end_date: 分析结束时间
    :param _filter_before_params: 过滤条件，与回放脚本写法相同
    :return: 过滤条件过滤后的该时间段内每个时刻的币种，分别为做多df1和做空df2
    """
    filter_before_exec = [filter_generate(param=param) for param in _filter_before_params]

    if filter_before_exec:
        print('前置过滤源码：')
        [print(x, '\n') if tag in x else print(x) for x in filter_before_exec]

    # 提取前置过滤因子 (in playback_start)
    filter_list = []
    [filter_list.extend(re.findall(r"\['(.+?)'\]", x)) for x in filter_before_exec]
    filter_list = list(set(filter_list))

    if 'fundingRate' in filter_list:
        use_fundingRate = True
        filter_list.remove('fundingRate')
    else:
        use_fundingRate = False
    start_date = pd.to_datetime(_start_date)
    end_date = pd.to_datetime(_end_date)

    filter_class_list = tools.convert_to_filter_cls(filter_list)

    # ===读取数据
    df = reader.readhour_filter(
        'swap',
        filter_class_list=filter_class_list)
    if df['candle_begin_time'].max() < pd.to_datetime(end_date):
        data_modify_time = df['candle_begin_time'].max() - timedelta(hours=1)
        log.warning(f'本地数据最新日期小于设定回测结束时间,请检查。本次回测结束时间将被改为:{data_modify_time}')
        end_date = data_modify_time
    if df['candle_begin_time'].min() > pd.to_datetime(start_date) - timedelta(hours=int(hold_hour_num)):
        data_modify_time = df['candle_begin_time'].min() + timedelta(hours=int(hold_hour_num))
        log.warning(f'本地数据最早日期大于设定回测开始时间,请检查。本次回测开始时间将被改为:{data_modify_time}')
        start_date = data_modify_time

    # 筛选日期范围
    df = df[df['candle_begin_time'] >= pd.to_datetime(
        start_date) - timedelta(hours=int(hold_hour_num))]
    df = df[df['candle_begin_time'] <= pd.to_datetime(end_date)]

    all_symbol_list = sorted(list(set(df['symbol'].unique())))
    replace_symbol_to_int = {v: k for k, v in enumerate(all_symbol_list)}  # symbol to int
    replace_symbol_to_int_ = {k: v for k, v in enumerate(all_symbol_list)}  # int to symbol
    df['symbol'] = df['symbol'].replace(replace_symbol_to_int)

    # 删除某些行数据
    df = df[df['volume'] > 0]  # 该周期不交易的币种
    # 最后几行数据,下个周期_avg_price为空
    df.dropna(subset=['下个周期_avg_price'], inplace=True)
    # ===数据预处理
    df = df[['candle_begin_time', 'close', 'symbol'] + filter_list]
    df = df.set_index(['candle_begin_time', 'symbol']).sort_index()
    df = df.replace([np.inf, -np.inf], np.nan)
    # 因子空值都用中位数填充, 如果填0可能后面rank排序在第一或者最后
    # df = df.fillna(value=0)
    df = df.reset_index()

    if use_fundingRate:
        # ===整合资金费率
        fundingrate_data = reader.read_fundingrate()
        fundingrate_data['symbol'] = fundingrate_data['symbol'].replace(replace_symbol_to_int)
        df = pd.merge(df,
                      fundingrate_data[['candle_begin_time', 'symbol', 'fundingRate']],
                      on=['candle_begin_time', 'symbol'], how="left")
        df['fundingRate'].fillna(value=0, inplace=True)
        log.info('整合资金费率完成')

    # 提前排除退市币种
    max_time = df['candle_begin_time'].max()
    quit_df = df.groupby('symbol')['candle_begin_time'].max().to_frame()
    quit_df = quit_df[quit_df['candle_begin_time'] < max_time]
    quit_symbols = quit_df.index.tolist()
    quit_df_ = df[df['symbol'].isin(quit_symbols)]
    noquit_df = df[~df['symbol'].isin(quit_symbols)]
    # 退市币种的处理,实盘提前N小时加入黑名单
    quit_df_ = quit_df_.groupby('symbol', group_keys=False).apply(
        lambda x: x.iloc[:-hold_hour_num - 1])
    df = noquit_df.append(quit_df_)
    df.sort_values(by=['candle_begin_time', 'symbol'], inplace=True)
    log.info('过滤数据处理完成')

    # np_gen_selected(function.py)函数前面部分的前置过滤
    df['candle_begin_time'] = pd.to_numeric(df['candle_begin_time'])
    base_time = pd.to_numeric(pd.Series(pd.to_datetime('20170101'))).iloc[0]
    df['offset'] = (df['candle_begin_time'] - base_time) / 3.6e12 % hold_hour_num

    # 前置过滤
    df1, df2 = filter_before(df.copy(), df.copy(), filter_before_exec, [long_white_list, short_white_list],
                             [long_black_list, short_black_list], replace_symbol_to_int)
    log.info('过滤筛选完成')

    df1 = df1[['candle_begin_time', 'symbol']]
    df2 = df2[['candle_begin_time', 'symbol']]
    df1['symbol'] = df1['symbol'].replace(replace_symbol_to_int_)
    df2['symbol'] = df2['symbol'].replace(replace_symbol_to_int_)
    df1['candle_begin_time'] = pd.to_datetime(df1['candle_begin_time'], format='%Y-%m-%d %H:%M:%S')
    df2['candle_begin_time'] = pd.to_datetime(df2['candle_begin_time'], format='%Y-%m-%d %H:%M:%S')

    return df1, df2


def get_all_symbol_df_with_factors(_start_date, _end_date, _factors_with_param, _df1, _df2):
    """
    根据过滤后的币种信息读取做多/空每个时刻每个币种的各个因子值，分别包含[candle_begin_time, symbol, factor1_param1, ...]多列
    :param _start_date: 分析开始时间(如果小于币种实际存在时间(未上架)，则使用币种实际存在时间)
    :param _end_date: 分析结束时间(如果大于币种实际存在时间(已下架)，则使用币种实际存在时间)
    :param _factors_with_param: 要读取的因子及参数，如[['MtmMean', 36], ['MtmMeanStd', 48], ...]
    :param _df1: 过滤后的做多币种信息
    :param _df2: 过滤后的做空币种信息
    :return: 做多/空的每个时刻每个币种的各个因子值
    """
    long_symbols = _df1[(_df1['candle_begin_time'] >= pd.to_datetime(_start_date)) & (
            _df1['candle_begin_time'] <= pd.to_datetime(_end_date))]
    short_symbols = _df2[(_df2['candle_begin_time'] >= pd.to_datetime(_start_date)) & (
            _df2['candle_begin_time'] <= pd.to_datetime(_end_date))]
    # 读取所有涉及币种的涨跌幅榜因子值
    all_symbols_after_filter = pd.unique(
        np.concatenate((long_symbols['symbol'].unique(), short_symbols['symbol'].unique()))).tolist()
    all_symbols_after_filter = [symbol.replace('USDT', '-USDT') for symbol in all_symbols_after_filter]
    all_long_symbol_df = []
    all_short_symbol_df = []
    for symbol in all_symbols_after_filter:
        symbol_df = pd.read_feather(os.path.join(pickle_path, symbol + '.pkl'))
        for _factor in _factors_with_param:
            factor_df = pd.read_feather(os.path.join(factor_pickle, symbol, 'coin_alpha_factor_' + _factor[0] + '.pkl'))
            symbol_df[f'{_factor[0]}_bh_{_factor[1]}'] = factor_df[f'{_factor[0]}_bh_{_factor[1]}']
        symbol_df = symbol_df[(symbol_df['candle_begin_time'] >= pd.to_datetime(_start_date)) & (
                symbol_df['candle_begin_time'] <= pd.to_datetime(_end_date))]
        _columns = ['candle_begin_time', 'symbol'] + [f'{_factor[0]}_bh_{_factor[1]}' for _factor in
                                                      _factors_with_param]
        symbol_df = symbol_df[_columns]  # 仅包含时间、币种名和多列因子值
        # 按照long_symbols和short_symbols筛选，仅保留每个时刻过滤后的有效币种，得到long_symbol_df和short_symbol_df
        long_symbol_df = pd.merge(symbol_df, long_symbols, on=['candle_begin_time', 'symbol'])
        short_symbol_df = pd.merge(symbol_df, short_symbols, on=['candle_begin_time', 'symbol'])
        all_long_symbol_df.append(long_symbol_df)
        all_short_symbol_df.append(short_symbol_df)
    all_long_symbol_df = pd.concat(all_long_symbol_df)
    all_short_symbol_df = pd.concat(all_short_symbol_df)
    return all_long_symbol_df, all_short_symbol_df


def get_and_output_rise_fall_symbols(_start_date, _end_date, _rise_fall_factor, _hold_period, _top_n,
                                     _filter_before_params, _save_path):
    """
    获取过滤后每个时刻做多和做空的所有币种中的前top_n个，并输出为csv文件
    :param _start_date: 分析开始时间(如果小于币种实际存在时间(未上架)，则使用币种实际存在时间)
    :param _end_date: 分析结束时间(如果大于币种实际存在时间(已下架)，则使用币种实际存在时间)
    :param _rise_fall_factor: 用于指示涨跌幅榜的因子
    :param _hold_period: rise_fall_factor因子对应参数
    :param _top_n: 要输出的币种中
    :param _filter_before_params: 过滤条件，与回放脚本写法相同
    :param _save_path: csv文件保存路径
    :return: 输出涨跌幅榜
    """
    _df1, _df2 = get_symbols_with_filter_before(_start_date, _end_date, _filter_before_params)
    factor_with_param = [[_rise_fall_factor, _hold_period]]
    all_long_symbol_df, all_short_symbol_df = get_all_symbol_df_with_factors(_start_date, _end_date, factor_with_param,
                                                                             _df1, _df2)

    # 分别对all_long_symbol_df和all_short_symbol_df，针对每个时刻对因子值'{rise_fall_factor}_bh_{hold_period}'进行排序
    all_long_symbol_df['long_rise_fall_rank'] = all_long_symbol_df.groupby('candle_begin_time')[
        f'{_rise_fall_factor}_bh_{_hold_period}'].rank(ascending=False)  # 降序，值最大排名1，过滤之后的涨幅榜
    all_short_symbol_df['short_rise_fall_rank'] = all_short_symbol_df.groupby('candle_begin_time')[
        f'{_rise_fall_factor}_bh_{_hold_period}'].rank(ascending=True)  # 升序，值最小排名1，过滤之后的跌幅榜

    # 获取每个时刻的排名前 top_n 的币种
    top_long_symbols = all_long_symbol_df[all_long_symbol_df['long_rise_fall_rank'] <= _top_n]
    top_short_symbols = all_short_symbol_df[all_short_symbol_df['short_rise_fall_rank'] <= _top_n]

    # 输出为 CSV 文件
    # 转换数据帧格式
    top_long_symbols_output = top_long_symbols.pivot_table(index='candle_begin_time', columns='long_rise_fall_rank',
                                                           values='symbol', aggfunc='first')
    top_short_symbols_output = top_short_symbols.pivot_table(index='candle_begin_time', columns='short_rise_fall_rank',
                                                             values='symbol', aggfunc='first')

    # 仅保留前 top_n 列
    top_long_symbols_output = top_long_symbols_output.iloc[:, :_top_n]
    top_short_symbols_output = top_short_symbols_output.iloc[:, :_top_n]
    # 重命名列名
    top_long_symbols_output.columns = range(1, _top_n + 1)
    top_short_symbols_output.columns = range(1, _top_n + 1)
    # 保存为CSV
    if not os.path.exists(_save_path):
        os.mkdir(_save_path)
    top_long_symbols_output.to_csv(os.path.join(_save_path, f'涨幅榜Top{_top_n}.csv'))
    top_short_symbols_output.to_csv(os.path.join(_save_path, f'跌幅榜Top{_top_n}.csv'))

    return top_long_symbols, top_short_symbols


def calculate_symbols_rise_fall_rank(_start_date, _end_date, _symbol_list, _rise_fall_factor, _hold_period, _df1, _df2):
    f"""
    计算一个币种列表(symbol_list)中，每个币种的涨跌幅榜(rise_fall_factor因子)排名
    :param _start_date: 分析开始时间(如果小于币种实际存在时间(未上架)，则使用币种实际存在时间)
    :param _end_date: 分析结束时间(如果大于币种实际存在时间(已下架)，则使用币种实际存在时间)
    :param _symbol_list: 要计算排名的币种名称，如['BTCUSDT', ...]
    :param _rise_fall_factor: 用于指示涨跌幅榜的因子
    :param _hold_period: rise_fall_factor因子对应参数
    :param _df1: 每个时刻过滤后的有效做多币种
    :param _df2: 每个时刻过滤后的有效做空币种
    :return: 包含每个币种的涨跌幅榜排名，包含[candle_begin_time, symbol, long_{_rise_fall_factor}_bh_{_hold_period}_rank, short_{_rise_fall_factor}_bh_{_hold_period}_rank]
    """
    _factor_with_param = [[_rise_fall_factor, _hold_period]]
    all_long_symbol_df, all_short_symbol_df = get_all_symbol_df_with_factors(_start_date, _end_date, _factor_with_param,
                                                                             _df1, _df2)

    # 分别对all_long_symbol_df和all_short_symbol_df，针对每个时刻对因子值'{rise_fall_factor}_bh_{hold_period}'进行排序
    all_long_symbol_df[f'long_{_rise_fall_factor}_bh_{_hold_period}_rank'] = all_long_symbol_df.groupby(
        'candle_begin_time')[f'{_rise_fall_factor}_bh_{_hold_period}'].rank(ascending=False)  # 降序，值最大排名1
    all_short_symbol_df[f'short_{_rise_fall_factor}_bh_{_hold_period}_rank'] = all_short_symbol_df.groupby(
        'candle_begin_time')[f'{_rise_fall_factor}_bh_{_hold_period}'].rank(ascending=False)  # 降序，值最大排名1

    # 分别对all_long_symbol_df和all_short_symbol_df，计算每个时刻symbol_list中每个币种的排名
    all_symbol_df = pd.merge(
        all_long_symbol_df[['candle_begin_time', 'symbol', f'long_{_rise_fall_factor}_bh_{_hold_period}_rank']],
        all_short_symbol_df[['candle_begin_time', 'symbol', f'short_{_rise_fall_factor}_bh_{_hold_period}_rank']],
        on=['candle_begin_time', 'symbol'],
        how='outer')

    all_symbol_df = all_symbol_df[all_symbol_df['symbol'].isin(_symbol_list)]
    all_symbol_df = all_symbol_df.sort_values(by=['candle_begin_time', 'symbol']).reset_index(drop=True)
    return all_symbol_df


def calculate_symbols_factor_list_rank(_start_date, _end_date, _symbol_list, _factor_list, _df1, _df2):
    """
    计算一个币种列表(symbol_list)中，每个币种的factor_list因子排名
    :param _start_date: 分析开始时间(如果小于币种实际存在时间(未上架)，则使用币种实际存在时间)
    :param _end_date: 分析结束时间(如果大于币种实际存在时间(已下架)，则使用币种实际存在时间)
    :param _symbol_list: 要计算排名的币种名称，如['BTCUSDT', ...]
    :param _factor_list: 要计算币种排名的因子列表，如[['MtmMean', 36], ['MtmMeanStd', 48], ...]
    :param _df1: 每个时刻过滤后的有效做多币种
    :param _df2: 每个时刻过滤后的有效做空币种
    :return: 包含每个币种的因子值排名，包含[candle_begin_time, symbol, long_factor1_bh_param1_rank, short_factor1_bh_param1_rank, ...]
    """
    all_long_symbol_df, all_short_symbol_df = get_all_symbol_df_with_factors(_start_date, _end_date, _factor_list,
                                                                             _df1, _df2)
    factors_column_name = [f'{factor[0]}_bh_{factor[1]}' for factor in _factor_list]

    # 分别对all_long_symbol_df和all_short_symbol_df，针对每个时刻对因子值进行排序
    for factor_column in factors_column_name:
        all_long_symbol_df[f'long_{factor_column}_rank'] = all_long_symbol_df.groupby('candle_begin_time')[
            factor_column].rank(ascending=False)  # 降序，值最大排名1
        all_short_symbol_df[f'short_{factor_column}_rank'] = all_long_symbol_df.groupby('candle_begin_time')[
            factor_column].rank(ascending=False)  # 降序，值最大排名1
    long_factors_column_name = [f'long_{factor_name}_rank' for factor_name in factors_column_name]
    short_factors_column_name = [f'short_{factor_name}_rank' for factor_name in factors_column_name]

    # 分别对all_long_symbol_df和all_short_symbol_df，计算每个时刻symbol_list中每个币种的排名
    # 如果该币种在该时刻不存在(被筛选掉了)，则设置为NAN；最终得到的df包含[candle_begin_time, symbol, long_rise_fall_rank, short_rise_fall_rank]
    all_symbol_df = pd.merge(all_long_symbol_df[['candle_begin_time', 'symbol'] + long_factors_column_name],
                             all_short_symbol_df[['candle_begin_time', 'symbol'] + short_factors_column_name],
                             on=['candle_begin_time', 'symbol'],
                             how='outer')

    all_symbol_df = all_symbol_df[all_symbol_df['symbol'].isin(_symbol_list)]
    all_symbol_df = all_symbol_df.sort_values(by=['candle_begin_time', 'symbol']).reset_index(drop=True)
    return all_symbol_df


def symbols_factors_analysis_and_draw(_start_date, _end_date, _symbol_list, _rise_fall_factor, _hold_period, _top_n,
                                      _factor_list, _filter_before_params, _save_path, if_draw_background=False):
    """
    对symbol_list中每个币种的k线、涨跌幅榜和因子列表进行画图
    :param _start_date: 分析开始时间(如果小于币种实际存在时间(未上架)，则使用币种实际存在时间)
    :param _end_date: 分析结束时间(如果大于币种实际存在时间(已下架)，则使用币种实际存在时间)
    :param _symbol_list: 要分析的币种列表，如['BTC', ...]
    :param _rise_fall_factor: 用于指示涨跌幅榜的
    :param _hold_period: 目标持有时长
    :param _top_n: 仅用于在排名折线图中画分割线
    :param _factor_list: 每个币种列表要分析的因子列表，如[['MtmMean', 36], ['MtmMeanStd', 48], ...]
    :param _filter_before_params: 过滤参数
    :param _save_path: 保存路径
    :param if_draw_background: 是否绘制背景(非常耗时)
    """
    # 计算过滤后的币种
    _df1, _df2 = get_symbols_with_filter_before(_start_date, _end_date, _filter_before_params)
    # 计算涨跌幅榜排名
    symbols_rise_fall_rank = calculate_symbols_rise_fall_rank(_start_date, _end_date, _symbol_list, _rise_fall_factor,
                                                              _hold_period, _df1, _df2)
    # 计算因子值排名
    symbols_factors_rank = calculate_symbols_factor_list_rank(_start_date, _end_date, _symbol_list, _factor_list, _df1,
                                                              _df2)

    for symbol in _symbol_list:
        print(f'绘制 {symbol} 的k线及因子曲线...')
        symbol_with_usdt = symbol + '-USDT'

        # 获取k线及因子数据
        draw_df = pd.read_feather(os.path.join(pickle_path, symbol_with_usdt + '.pkl'))
        for i, _factor in enumerate(_factor_list):
            factor_df = pd.read_feather(
                os.path.join(factor_pickle, symbol_with_usdt, 'coin_alpha_factor_' + f'{_factor[0]}' + '.pkl'))
            draw_df[f'{_factor[0]}_bh_{_factor[1]}'] = factor_df[f'{_factor[0]}_bh_{_factor[1]}']
        draw_df = draw_df[(draw_df['candle_begin_time'] >= pd.to_datetime(_start_date)) &
                          (draw_df['candle_begin_time'] <= pd.to_datetime(_end_date))]

        # 获取涨跌幅榜的排名数据(如果不为None)
        symbol_usdt = symbol + 'USDT'
        rise_fall_rank = symbols_rise_fall_rank[symbols_rise_fall_rank['symbol'] == symbol_usdt]
        # 获取各因子值的排名数据(如果不为None)
        factors_rank = symbols_factors_rank[symbols_factors_rank['symbol'] == symbol_usdt]
        rank_df = pd.merge(rise_fall_rank, factors_rank, on='candle_begin_time', how='outer')
        draw_df = pd.merge(draw_df, rank_df, on='candle_begin_time', how='outer')

        # 计算每个时刻过滤后多/空总币种数，用于显示多空涨跌幅榜/因子排名时，按照"pos_rank/total-pos_rank"的方式显示
        long_symbol_count = _df1.groupby('candle_begin_time')['symbol'].nunique().reset_index(name='long_symbol_num')
        short_symbol_count = _df2.groupby('candle_begin_time')['symbol'].nunique().reset_index(name='short_symbol_num')
        # 合并到一个新的 DataFrame
        symbol_count = pd.merge(long_symbol_count, short_symbol_count, on='candle_begin_time', how='outer')
        draw_df = pd.merge(draw_df, symbol_count, on='candle_begin_time', how='outer')

        # 创建子图
        fig = make_subplots(rows=3,
                            cols=1,
                            shared_xaxes=True,
                            vertical_spacing=0.02,
                            specs=[[{"secondary_y": True}], [{}], [{}]],  # 添加次级Y轴
                            row_heights=[4, 1.2, 0.8])

        # 设置图表布局
        fig.update_layout(
            hovermode='x unified',  # 设置统一的悬停模式
            hoverlabel=dict(font_size=10),  # 悬停字体大小
        )

        # 添加 K 线图
        candlestick = go.Candlestick(x=draw_df['candle_begin_time'],
                                     open=draw_df['open'],
                                     high=draw_df['high'],
                                     low=draw_df['low'],
                                     close=draw_df['close'],
                                     name=f'{symbol}_kline',
                                     increasing_line_color='green',
                                     decreasing_line_color='red')
        fig.add_trace(candlestick, row=1, col=1, secondary_y=False)

        # 绘制多个因子值
        for i, _factor in enumerate(_factor_list):
            factor_trace = go.Scatter(
                x=draw_df['candle_begin_time'],
                y=draw_df[f'{_factor[0]}_bh_{_factor[1]}'],
                mode='lines',
                name=f'因子值：{_factor[0]}_bh_{_factor[1]}')
            fig.add_trace(factor_trace, row=1, col=1, secondary_y=True)

        # 添加涨跌幅排名和因子值排名的折线图
        _factor_list_with_rff = [[_rise_fall_factor, _hold_period]] + _factor_list
        for i, _factor in enumerate(_factor_list_with_rff):
            long_rank_trace = go.Scatter(
                x=draw_df['candle_begin_time'],
                y=draw_df[f'long_{_factor[0]}_bh_{_factor[1]}_rank'],
                mode='lines+markers',
                name=f'long_{_factor[0]}_bh_{_factor[1]}_rank',
            )
            fig.add_trace(long_rank_trace, row=2, col=1)

            short_rank_trace = go.Scatter(
                x=draw_df['candle_begin_time'],
                y=draw_df[f'short_{_factor[0]}_bh_{_factor[1]}_rank'],
                mode='lines+markers',
                name=f'short_{_factor[0]}_bh_{_factor[1]}_rank',
            )
            fig.add_trace(short_rank_trace, row=2, col=1)

        # 在折线图中添加总数
        long_num_trace = go.Scatter(
            x=draw_df['candle_begin_time'],
            y=draw_df['long_symbol_num'],
            mode='lines+markers',
            name='long_total',
        )
        fig.add_trace(long_num_trace, row=2, col=1)

        short_num_trace = go.Scatter(
            x=draw_df['candle_begin_time'],
            y=draw_df['short_symbol_num'],
            mode='lines+markers',
            name='short_total',
        )
        fig.add_trace(short_num_trace, row=2, col=1)

        # 在折线图中添加top_n分割线
        draw_df['top_n'] = _top_n
        long_num_trace = go.Scatter(
            x=draw_df['candle_begin_time'],
            y=draw_df['top_n'],
            mode='lines',
            name='long_top_n',
            showlegend=False,  # 不显示在图例中
        )
        fig.add_trace(long_num_trace, row=2, col=1)

        short_num_trace = go.Scatter(
            x=draw_df['candle_begin_time'],
            y=draw_df['short_symbol_num'] - _top_n,
            mode='lines',
            name='short_top_n',
            showlegend=False,  # 不显示在图例中
        )
        fig.add_trace(short_num_trace, row=2, col=1)

        # 添加成交额图
        volume_bar = go.Bar(x=draw_df['candle_begin_time'],
                            y=draw_df['quote_volume'],
                            name='成交额',
                            marker_color='green'
                            )
        fig.add_trace(volume_bar, row=3, col=1)

        if if_draw_background:
            # 添加涨跌幅榜内背景(浅蓝)
            condition_met = (draw_df[f'long_{_rise_fall_factor}_bh_{_hold_period}_rank'] <= _top_n) | \
                            (draw_df[f'short_{_rise_fall_factor}_bh_{_hold_period}_rank'] >= draw_df[
                                'short_symbol_num'] - _top_n)
            # 标记连续满足条件的时间段
            draw_df['condition_met'] = condition_met
            draw_df['group'] = (draw_df['condition_met'] != draw_df['condition_met'].shift()).cumsum()

            for _, group in draw_df[draw_df['condition_met']].groupby('group'):
                start_time = group['candle_begin_time'].min() - timedelta(minutes=30)
                end_time = group['candle_begin_time'].max() + timedelta(minutes=30)

                # 绘制浅绿色背景
                fig.add_vrect(x0=start_time, x1=end_time, fillcolor="LightGreen", opacity=0.4, layer="below",
                              line_width=0)

        # 设置图表布局
        fig.update_layout(
            title=f'{symbol}_factor_analysis',
            xaxis_rangeslider_visible=False,
            font=dict(size=16, color='blue'),
        )

        # 设置title
        fig.update_yaxes(title_text="K线", row=1, col=1)
        fig.update_yaxes(title_text="因子值", row=1, col=1, secondary_y=True)
        fig.update_yaxes(title_text="因子排名", row=2, col=1, autorange='reversed')
        fig.update_yaxes(title_text="成交额", row=3, col=1)

        # 保存图片
        if not os.path.exists(_save_path):
            os.mkdir(_save_path)
        start_symbol = datetime.strptime(_start_date, "%Y-%m-%d").strftime("%Y%m%d")
        end_symbol = datetime.strptime(_end_date, "%Y-%m-%d").strftime("%Y%m%d")
        fig.write_html(os.path.join(_save_path, f'{symbol}_{start_symbol}_{end_symbol}.html'))

